{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ffdb6fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from atomworks.io import parse\n",
    "from atomwork.io.tools.inference import components_to_atom_array\n",
    "from atomworks.io.utils.visualize import view\n",
    "from atomworks.io.utils.io_utils import to_cif_file\n",
    "\n",
    "import os\n",
    "\n",
    "CURRENT_DIR = os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20cb0d36",
   "metadata": {},
   "outputs": [],
   "source": [
    "protein = {\n",
    "    \"seq\": \"MTPEERAERAREYADYHVDYVGERVGLPVTPEMRLFDELVYRYMVARDPSILPELERYRELGDRLAARRADPRWAGKSRAELYEEFLTPEEREEVREFNRRYAEAVRENRERLAREVGVPEEVMDRVIAAMDEAYAAGESV(L:0)EAYYRVLHEAVAERAA\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47b03b36",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_residues = {\n",
    "    \"L:0\": {\n",
    "        \"path\": f\"{CURRENT_DIR}/penicillin_ts2_as_ncaa.cif\",\n",
    "        \"chain_type\": \"non-polymer\"\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67ac2b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "view(parse(custom_residues[\"L:0\"][\"path\"], hydrogen_policy=\"remove\")[\"asym_unit\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32cc5df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "atom_array = components_to_atom_array([protein], custom_residues=custom_residues)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "425e57d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the bonds\n",
    "\n",
    "# TODO: CHECK BONDS\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "707c34ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_cif_file(atom_array, f\"{CURRENT_DIR}/ligand_as_ncaa.cif\", include_entity_poly=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07a281c4",
   "metadata": {},
   "source": [
    "# Ensure that our outputs have the right chirality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f622a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from foundry.metrics.chiral import compute_chiral_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d96bdd27",
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted_atom_array = parse(\"/mnt/home/ncorley/ligand_as_ncaa.cif.gz\")[\"asym_unit\"][0]\n",
    "\n",
    "# ... subset to our ligand, L:0\n",
    "predicted_ncaa_atom_array = predicted_atom_array[predicted_atom_array.res_name == \"L:0\"]\n",
    "\n",
    "# (Ground truth)\n",
    "ground_truth_ncaa_atom_array = parse(\"/home/ncorley/projects/modelhub_dev/tests/data/ncaa/penicillin_ts2_as_ncaa.cif\", hydrogen_policy=\"remove\")[\"asym_unit\"][0]\n",
    "\n",
    "assert len(predicted_ncaa_atom_array) == len(ground_truth_ncaa_atom_array), \"Number of atoms in predicted and ground truth NCAA do not match!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89fdd327",
   "metadata": {},
   "outputs": [],
   "source": [
    "view(predicted_ncaa_atom_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1ad9b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "view(ground_truth_ncaa_atom_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d8d1a00",
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_chiral_metrics(\n",
    "    predicted_ncaa_atom_array,\n",
    "    ground_truth_ncaa_atom_array\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "modelhub",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
